package com.cmw.socket.impl;

import com.cmw.socket.core.IoProvider;
import com.cmw.socket.utils.CloseUtils;

import java.io.IOException;
import java.nio.channels.ClosedChannelException;
import java.nio.channels.SelectionKey;
import java.nio.channels.Selector;
import java.nio.channels.SocketChannel;
import java.util.HashMap;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * @ClassName IoSelectorProvider
 * @Description: IO Selector 选择器提供对象
 * @Author cheng
 * @Date 2020/6/24 10:05
 * @Version V1.0
 **/
public class IoSelectorProvider implements IoProvider {

    private final AtomicBoolean isClosed = new AtomicBoolean(false);

    //是否处于某个过程
    private final AtomicBoolean inRegInput = new AtomicBoolean(false);
    private final AtomicBoolean inRegOutput = new AtomicBoolean(false);

    private final Selector readSelector;
    private final Selector writeSelector;

    private final Map<SelectionKey, Runnable> inputCallbackMap = new HashMap<>();
    private final Map<SelectionKey, Runnable> outputCallbackMap = new HashMap<>();

    private final ExecutorService inputHandlePool;
    private final ExecutorService outputHandlePool;

    public IoSelectorProvider() throws IOException {
        readSelector = Selector.open();
        writeSelector = Selector.open();
        inputHandlePool = Executors.newFixedThreadPool(4, new IoProviderThreadFactory("IoProvider-Input-Thread-"));
        outputHandlePool = Executors.newFixedThreadPool(4, new IoProviderThreadFactory("IoProvider-Output-Thread-"));
        //开始输出输入的监听
        startRead();
        startWrite();
    }

    /**
     * 开始读线程
     */
    private void startRead(){
        Thread thread = getThreadByReadOrWrite(1);
    }

    /**
     * 开始写线程
     */
    private void startWrite(){
        Thread thread = getThreadByReadOrWrite(2);
    }

    private Thread getThreadByReadOrWrite(int ioType) {
        String threadName = null;
        Map<SelectionKey, Runnable> callbackMap = null;
        ExecutorService handlePool = null;
        AtomicBoolean inRegIoPut = null;
        Selector ioSelector = null;

        int opIoType = 0;
        if(ioType == 1){ //读
            threadName = "Clink IoSelectorProvider ReadSelector Thread";
            ioSelector = readSelector;
            callbackMap = inputCallbackMap;
            handlePool = inputHandlePool;
            inRegIoPut = inRegInput;
            opIoType = SelectionKey.OP_READ;
        }else{//写
            threadName = "Clink IoSelectorProvider WriteSelector Thread";
            ioSelector = writeSelector;
            callbackMap = outputCallbackMap;
            handlePool = outputHandlePool;
            inRegIoPut = inRegOutput;
            opIoType = SelectionKey.OP_WRITE;
        }
        final Selector _ioSelector = ioSelector;
        final Map<SelectionKey, Runnable> _callbackMap = callbackMap;
        final ExecutorService _handlePool = handlePool;
        final AtomicBoolean _inRegIoPut = inRegIoPut;
        final int _opIoType = opIoType;
        Thread thread = new Thread(threadName){
            @Override
            public void run() {
                while (!isClosed.get()){
                    try {
                        if(_ioSelector.select() == 0){
                            waitSelection(_inRegIoPut);
                            continue;
                        }
                        Set<SelectionKey> selectionKeys = _ioSelector.selectedKeys();
                        for(SelectionKey selectionKey : selectionKeys){
                            if(selectionKey.isValid()){
                                handleSelection(selectionKey, _opIoType, _callbackMap, _handlePool);
                            }
                        }
                        selectionKeys.clear();
                    } catch (IOException e) {
                        e.printStackTrace();
                    }
                }
            }
        };
        thread.setPriority(Thread.MAX_PRIORITY);
        thread.start();
        return thread;
    }

    private void handleSelection(SelectionKey key, int keyOps, Map<SelectionKey, Runnable> map, ExecutorService handlePool) {
        //重点
        //取消继续对keyOps的监听
        key.interestOps(key.readyOps() & ~keyOps);
        Runnable runnable = null;
        try{
            runnable = map.get(key);
        }catch (Exception ignored){

        }
        if(null != runnable && !handlePool.isShutdown()){
            //异步调度
            handlePool.execute(runnable);
        }
    }


    private static void waitSelection(final AtomicBoolean locker){
        //noinspection SynchronizationOnLocalVariableOrMethodParameter
        synchronized (locker){
            if(locker.get()){
                try {
                    locker.wait();
                } catch (InterruptedException e) {
                    e.printStackTrace();
                }
            }
        }
    }

    @Override
    public boolean registerInput(SocketChannel channel, HandleInputCallback callback) {
        return registerSelection(channel, readSelector, SelectionKey.OP_READ, inRegInput, inputCallbackMap, callback) != null;
    }

    private SelectionKey registerSelection(SocketChannel channel, Selector selector, int registerOps, AtomicBoolean locker, Map<SelectionKey, Runnable> map, Runnable runnable) {
        //noinspection SynchronizationOnLocalVariableOrMethodParameter
        synchronized (locker){
            // 设置锁定状态
            locker.set(true);
            try {
                // 唤醒当前的selector，让selector不处于select()状态
                selector.wakeup();
                SelectionKey key = null;
                if(channel.isRegistered()){
                    //查询是否已经注册过
                    key = channel.keyFor(selector);
                    if(null != key){
                        key.interestOps(key.readyOps() | registerOps);
                    }
                }
                if(null == key){
                    // 注册selector得到Key
                    key = channel.register(selector, registerOps);
                    //注册回调
                    map.put(key, runnable);
                }
                return key;
            } catch (ClosedChannelException e) {
                e.printStackTrace();
                return null;
            }finally {
                //解除锁定状态
                locker.set(false);
                try{
                    //通知
                    locker.notify();
                }catch (Exception ignored){
                    ignored.printStackTrace();
                }
            }
        }
    }

    @Override
    public boolean registerOutput(SocketChannel channel, HandleOutputCallback callback) {
        return registerSelection(channel, writeSelector, SelectionKey.OP_WRITE, inRegOutput, outputCallbackMap, callback) != null;
    }

    @Override
    public void unRegisterInput(SocketChannel channel) {
        unRegisterSelection(channel, readSelector, inputCallbackMap);
    }

    @Override
    public void unRegisterOutput(SocketChannel channel) {
        unRegisterSelection(channel, writeSelector, outputCallbackMap);
    }

    private void unRegisterSelection(SocketChannel channel, Selector selector, Map<SelectionKey, Runnable> map){
        if(channel.isRegistered()){
            SelectionKey key = channel.keyFor(selector);
            if(null != key){
                //取消监听的方法
                key.cancel();
                map.remove(key);
                selector.wakeup();  // 唤醒当前的selector，让selector不处于select()状态
            }
        }
    }

    @Override
    public void close() throws IOException {
        if(isClosed.compareAndSet(false, true)){
            inputHandlePool.shutdown();
            outputHandlePool.shutdown();
            inputCallbackMap.clear();
            outputCallbackMap.clear();
            readSelector.wakeup();
            writeSelector.wakeup();
            CloseUtils.close(readSelector, writeSelector);
        }
    }

    static class IoProviderThreadFactory implements ThreadFactory{
        private final ThreadGroup group;
        private final AtomicInteger threadNumber = new AtomicInteger(1);
        private final String namePrefix;

        public IoProviderThreadFactory(String namePrefix) {
            SecurityManager s = System.getSecurityManager();
            this.group = (null != s) ? s.getThreadGroup() : Thread.currentThread().getThreadGroup();
            this.namePrefix = namePrefix;
        }

        @Override
        public Thread newThread(Runnable r) {
            Thread t = new Thread(group, r, namePrefix+threadNumber.getAndIncrement(), 0);
            if(t.isDaemon()) t.setDaemon(false);
            if(t.getPriority() != Thread.NORM_PRIORITY) t.setPriority(Thread.NORM_PRIORITY);
            return t;
        }
    }
}
